#!/usr/bin/python
import os
import astro.spec
import matplotlib.pyplot as pl
from astro.utilities import readatom, findtrans, indexnear,calc_abs
#from astro.findline import find_feature_centre
from matplotlib.ticker import NullLocator
import numpy as np
from math import sqrt

def to_vpfitname(name):
    isplit = 2 if name[1] not in 'IVX' else 1
    atomname = name[:isplit]
    ionisation = name[isplit:] 
    return '%s %s' % (atomname, ionisation)

# read in one or more spectra.  left and right arrows to switch between them

usage = 'Usage:  plotvel.py filename [redshift]'
help = """
space        Mark a new redshift position and redraw.
e            Measure the observed equivalent width over a region. Place the
             cursor on one region edge and press 'e', then repeat for the
             other edge.
"""

atom = readatom('/home/nhmc/installed/vpfit/atom.dat')
wlya = 1215.6701

# few strong transitions.  Could maybe add OVI doublet, OI 1302, FeII,
# AlII, AlIII, MgII doublet, CaII?
filename = os.path.join(os.path.expanduser('~'), '.astropy/transitions.dat')
if os.path.lexists(filename):
    fh = open(filename)
    TRANS = fh.readlines()
else:
    TRANS = ('CIV 1550, CIV 1548,  SiIV 1403, SiIV 1393,'
             'CII 1334, SiII 1260, HI 1215,   SiIII 1206')

Ckms = 299792.458         # speed of light km/s, exact

def initvelplot(wa, nfl, ner, nco, transitions, z, fig,
            vmin=-1000., vmax=1000., ticks=None):
    """ Vertical stacked plots of expected positions of absorption
    lines at the given redshift on a velocity scale."""

    colours = ('b','r', 'g', 'y', 'c', 'purple')
    ions = [tr[0].split()[0] for tr in transitions]
    # want an ordered set
    ionset = []
    for ion in ions:
        if ion not in ionset:
            ionset.append(ion)

    colour = dict(zip(ionset, colours * (len(ions) // len(colours) + 1)))

    n = len(transitions)
    zp1 = z + 1
    betamin = vmin / Ckms
    betamax = vmax / Ckms

    #fig.subplots_adjust(hspace=0.001)
    ax = fig.add_axes([0.05,0.06,0.9,0.9])
    ax.set_autoscale_on(0)
    # plot top down, so we need reversed()
    offsets = []
    artists = dict(fl=[], er=[], co=[], tick=[], text=[])
    for i,(tstring,trans) in enumerate(reversed(transitions)):
        ion = tstring.split()[0]
        offset = i*1.5
        offsets.append(offset)
        watrans = trans['wav']
        obswa = watrans * zp1
        wmin = obswa * (1 + betamin)
        wmax = obswa * (1 + betamax)

        if ticks is not None:
            cond = (wmin < ticks) & (ticks < wmax)
            if np.any(cond):
                vel = (ticks[cond] / obswa - 1) * Ckms
                a = ax.vlines(vel, 1.1+offset, 1.35+offset, colors='k')
                artists['tick'].append(a)
                
        cond = (wmin < wa) & (wa < wmax)
        #good = ~np.isnan(fl) & (er > 0) & ~np.isnan(co)
        
        fl = nfl[cond]
        er = ner[cond]
        co = nco[cond]
        vel = (wa[cond] / obswa - 1) * Ckms
        ax.axhline(offset, color='gray', alpha=0.7)
        artists['er'].extend( 
            ax.plot(vel, er + offset, lw=1, color='orange', alpha=0.5) )
        artists['fl'].extend(
            ax.plot(vel, fl + offset, colour[ion], lw=1, alpha=0.7,
                    ls='steps-mid') )
        artists['co'].extend(
            ax.plot(vel, co + offset, color='gray', alpha=0.7, ls='dashed'))

        bbox = dict(facecolor='w', edgecolor='None')
        text = ax.text(0.9*vmin, offset+0.5,tstring, fontsize=15, bbox=bbox)
        artists['text'].append(text)
    artists['vline'] = ax.axvline(0, color='k', linestyle='dotted')
    ax.set_xlim(vmin, vmax)
    ax.set_ylim(-0.5, len(transitions) * 1.5)
    ax.yaxis.set_major_locator(NullLocator())
    ax.set_xlabel('Velocity offset (km s$^{-1}$)',fontsize=16)
    ax.set_title('$z = %.5f$' % z, fontsize=18)

    return artists, np.array(offsets)


class VelplotWrap(object):
    def __init__(self, wa, nfl, ner, fig, z, filename,
                 transitions=TRANS, vmin=-1000., vmax=1000.):
        if isinstance(transitions, basestring):
            transitions = transitions.split(',')
        self.trans = [findtrans(tr, atomdat=atom) for tr in transitions]
        self.filename = filename
        self.cids = dict()
        # disable any existing key press callbacks
        cids = list(fig.canvas.callbacks.callbacks['key_press_event'])
        for cid in cids:
            fig.canvas.callbacks.disconnect(cid)
        self.connect(fig)
        self.wa = wa
        self.ner = ner
        self.nfl = nfl
        self.co = np.ones(len(wa))
        edges = astro.spec.find_wa_edges(wa)    # pixel edges
        dw = edges[1:] - edges[:-1]             # width per pixel
        self.ew =  (1 - nfl) * dw               # equiv. width
        self.ewer = dw * ner
        self.fig = fig
        self.ticks = None
        self.vmin = vmin
        self.vmax = vmax
        self.prev_keypress = None
        self.prev_wa = None
        self.indprev = None
        artists,offsets = initvelplot(
            wa, nfl, ner, self.co, self.trans, z, fig,
            vmin=vmin, vmax=vmax, ticks=self.ticks)
        self.offsets = offsets
        self.artists = artists
        self.update(z)

    def update(self, z):
        wa, nfl, ner, nco = self.wa, self.nfl, self.ner, self.co
        self.z = z

        ax = pl.gca()
        zp1 = z + 1
        betamin = self.vmin / Ckms
        betamax = self.vmax / Ckms
        artists = self.artists
        ticks = self.ticks
        for t in artists['tick']:
            t.remove()
        artists['tick'] = []
        
        # want plots to appear from  top down, so we need reversed()
        for i,(tstring,trans) in enumerate(reversed(self.trans)):
            offset = i*1.5
            watrans = trans['wav']
            obswa = watrans * zp1
            wmin = obswa * (1 + betamin)
            wmax = obswa * (1 + betamax)

            if ticks is not None:
                cond = (wmin < ticks) & (ticks < wmax)
                if np.any(cond):
                    vel = (ticks[cond] / obswa - 1) * Ckms
                    a = ax.vlines(vel, 1.1+offset, 1.35+offset, colors='k')
                    artists['tick'].append(a)

            cond = (wmin < wa) & (wa < wmax)
            #good = ~np.isnan(fl) & (er > 0) & ~np.isnan(co)
        
            fl = nfl[cond]
            er = ner[cond]
            co = nco[cond]
            vel = (wa[cond] / obswa - 1) * Ckms
            artists['fl'][i].set_xdata(vel)
            artists['fl'][i].set_ydata(fl + offset)
            artists['er'][i].set_xdata(vel)
            artists['er'][i].set_ydata(er + offset)
            artists['co'][i].set_xdata(vel)
            artists['co'][i].set_ydata(co + offset)

            artists['text'][i].set_x(0.9 * self.vmin)

        ax.set_xlim(self.vmin, self.vmax)
        ax.set_ylim(-0.5, len(self.trans) * 1.5)
        ax.set_title('$z = %.5f$' % self.z, fontsize=18)
        pl.draw()

        self.artists = artists


    def on_keypress_navigate(self, event):
        """ move up and down the velocity spectrum
        """
        if event.inaxes is None:
            return
        z = self.z
        if event.key == ' ':
            # get new redshift
            dz = (event.xdata / Ckms) * (1 + z)
            # redraw()
            #print 'new redshift', z+dz
            self.vmin, self.vmax = pl.gca().get_xlim()
            self.update(z + dz)            
            
    def on_keypress_print(self,event):
        """ Print f26-style fitting regions and line guesses.
        """
        if event.key in 'ble':
            i = self.offsets.searchsorted(event.ydata)
            if i == 0:
                i = len(self.trans)
            name, tr = self.trans[-i]
            z = (1 + event.xdata / Ckms) * (1 + self.z) - 1
            wa = tr['wav'] * (1 + z)
            ind = indexnear(self.wa, wa)
            print 'transition', name, tr, wa
        if event.key == 'b':
            # print fitting region
            if self.prev_keypress == 'b':
                wmin = self.prev_wa
                wmax = wa
                if wmin > wmax:
                    wmin, wmax = wmax, wmin
                print wmin, wmax
                print '%%%% %s 1 %.2f %.2f' % (self.filename, wmin, wmax)
                self.prev_keypress = None
            else:
                self.prev_wa = wa
        elif event.key == 'l':
            # guess a line N, z and b and print to screen need to know
            # which transition we want, and what the central index is.
            print '%s %.6f, %.2f, %.4f' % (to_vpfitname(name), 10, 13.0)
        elif event.key == 'e':
            if self.prev_keypress == 'e':
                i0,i1 = self.indprev, ind
                if i0 > i1:
                    i0, i1 = i1, i0
                ew = self.ew[i0:i1+1]
                ewer = self.ewer[i0:i1+1]
                good = ~np.isnan(ew) & (ewer > 0)
                ewtot = sum(ew[good])
                ewtoter = sqrt(sum(ewer[good]**2))
                #wac = find_feature_centre(self.wa[i0:i1+1], ew)
                #z = wac / tr['wav']  - 1
                wmin,wmax = self.wa[i0], self.wa[i1]
                print '%s %.7f %03i pix %.4f %.4f %.5f %.5f' % (
                    name, self.z, len(ewer[good]), ewtot, ewtoter, wmin, wmax)
                self.prev_keypress = None
                return
            else:
                self.indprev = ind

        self.prev_keypress = event.key

    def on_keypress_plotline(self, event):
        """ generate a voigt profile for a (lya) line
        """
        if event.inaxes is None:
            return
        if event.key == 'L':
            zp1 = 1 + self.z
            z = self.z + (event.xdata / Ckms * zp1) 
            wmin = (1 - 30000 / Ckms) * zp1 * wlya
            wmax = (1 + 30000 / Ckms) * zp1 * wlya
            i,j = self.wa.searchsorted( (wmin, wmax) )
            logN = raw_input('log column density? (21)')
            try:
                logN = float(logN)
            except ValueError:
                logN = 21
            co = self.co[i:j]
            wa = self.wa[i:j]
            co = calc_abs(wa, atom['HI'], atom['HI'][0], z+1,
                          30000, logN=logN, b=20.)
            self.co[i:j] = co
            self.wa[i:j] = wa
            print 'updating with new continuum'
            self.update(self.z)
        elif event.key == 'R':
            # reset continuum
            self.co = np.ones(len(self.wa))
            self.update(self.z)
        
    def connect(self, fig):
        cids = dict()
        cids['nav'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_navigate)
        cids['print'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_print)
        cids['plotline'] = fig.canvas.mpl_connect(
            'key_press_event', self.on_keypress_plotline)
        self.cids.update(cids)

#     def disconnect(self):
#         for key in self.cids:
#             self.fig.canvas.mpl_disconnect(self.cids[key])

def main(args):
    Nargs = len(args)
    z =  float(args[1]) if Nargs > 1 else 2.47
    dv =  float(args[2]) if Nargs > 2 else 1000

    filename = os.path.abspath(args[0])
    spec = astro.spec.read(filename)
    fig = pl.figure(figsize=(14,11))
    print help
    junk = VelplotWrap(spec.wa, spec.fl/spec.co, spec.er/spec.co, fig, z,
                       spec.filename, vmin=-dv, vmax=dv)
    pl.show()

if __name__ == '__main__':
    import sys
    sys.exit(main(sys.argv[1:]))
